import torch
import torch.nn as nn
from torch.cuda.amp import GradScaler, autocast
from torchattacks.attack import Attack
from attack.libfastattack.lossfunc import *

class FastPGD(Attack):

    def __init__(self, model, cls_num_list=None, eps=0.3,
                 alpha=2/255, steps=40, random_start=True, mode="Normal",weights=False):
        super().__init__("FastPGD", model)
        self.eps = eps
        self.alpha = alpha
        self.steps = steps
        self.random_start = random_start
        self._supported_mode = ['default', 'targeted']
        self.scaler = GradScaler()
        self.cls_num_list = cls_num_list
        self.mode = mode
        self.weights = weights

    def forward(self, images, labels):
        r"""
        Overridden.
        """
        mode = self.mode
        images = images.clone().detach().to(self.device)
        labels = labels.clone().detach().to(self.device)

        if self._targeted:
            target_labels = self._get_target_label(images, labels)
        if mode == "Normal":
            loss = nn.CrossEntropyLoss()
        elif mode == "CalFAT":
            loss = nn.KLDivLoss(size_average=False)
            marloss = maxMarginLoss_kl(cls_num_list=self.cls_num_list, s=10, weight=None).cuda()
        else:
            assert 1 == 0
        adv_images = images.clone().detach()
        Kappa = torch.ones(len(images)).cuda()
        

        if self.random_start:
            # Starting at a uniformly random point
            adv_images = adv_images + torch.empty_like(adv_images).uniform_(-self.eps, self.eps)
            adv_images = torch.clamp(adv_images, min=0, max=1).detach()

        for _ in range(self.steps):
            adv_images.requires_grad = True

            # Accelerating forward propagation
            with autocast():
                outputs = self.model(adv_images)
                origin_outputs = self.model(images)
                # Calculate loss
                if mode == "Normal":
                    if self._targeted:
                        cost = -loss(outputs, target_labels)
                    else:
                        cost = loss(outputs, labels)
                elif mode == "CalFAT":
                    cost = loss(
                        F.log_softmax(marloss(outputs), dim=1),
                        F.softmax(marloss(origin_outputs), dim =1)
                    )
                else:
                    assert 1 == 2
            # Update adversarial images with gradient scaler applied
            scaled_loss = self.scaler.scale(cost)
            grad = torch.autograd.grad(scaled_loss, adv_images,
                                       retain_graph=False, create_graph=False)[0]

            adv_images = adv_images.detach() + self.alpha*grad.sign()
            delta = torch.clamp(adv_images - images, min=-self.eps, max=self.eps)
            adv_images = torch.clamp(images + delta, min=0, max=1).detach()
            bad_output = F.softmax(outputs.detach())
            predict = torch.argmax(bad_output, dim=1)
            correct_predictions = (predict == labels).float()
            Kappa += correct_predictions
        if self.weights:
            return adv_images,Kappa
        else:
            return adv_images